#include <jni.h>
#include <string>
#include <JavaManager.h>
#include <LlmManager.h>

#define SIMPLE 0
// 假设 JVM 和回调对象是全局变量
JavaVM* g_javaVM = nullptr; //全局JavaVM 指针

void callback(RKLLMResult *result, void *userdata, LLMCallState state){
    static JNIEnv* env = nullptr;
    if(env == nullptr)
    {
        // 将当前线程附加到 JVM
        if (g_javaVM->AttachCurrentThread(&env, nullptr) != JNI_OK) {
            LOGE("Failed to attach thread to JVM\n");
            return;
        }
    }
    switch (state) {
        case RKLLM_RUN_FINISH:
            LOGI("run finsh\n");
#if SIMPLE
            JavaSimpleManager::callJavaFinish(env);
#else
            JavaManager::callJavaFinish(env);
#endif
            break;

        case RKLLM_RUN_ERROR:
            LOGE("run error\n");
            break;

        case RKLLM_RUN_GET_LAST_HIDDEN_LAYER:
            break;

        case RKLLM_RUN_NORMAL:
            LOGI("result is %s\n", result->text);
#if SIMPLE
            JavaSimpleManager::callJavaSimple(env, result->text, state);
#else
            JavaManager::callLlmResult(env, result, userdata, state);
#endif
            break;

        default:
            break;
    }

}

extern "C" JNIEXPORT jboolean  JNICALL
Java_com_taiic_deepseek_DsManager_initialize(
        JNIEnv* env,
        jobject /* this */,
        jstring modelPath, jint maxNewTokens, jint maxContextLen) {

    //    保存JavaVM指针,JavaVM是全局的，可以在任意线程中调用
    if(env->GetJavaVM((&g_javaVM)) != JNI_OK)
    {
        LOGE("Failed to get JavaVM");
        return JNI_FALSE;
    }

#if SIMPLE
    JavaSimpleManager::initJavaSimple(env);
#else
    JavaManager::initJavaMethod(env);
#endif

    // 将 jstring 转换为 C++ 字符串
    const char* nativeModelPath = env->GetStringUTFChars(modelPath, nullptr);
    if (nativeModelPath == nullptr) {
        // 内存分配失败
        return JNI_FALSE;
    }
// 打印参数（仅用于调试）
    LOGD("Model Path: %s", nativeModelPath);
    LOGD("Max New Tokens: %d", maxNewTokens);
    LOGD("Max Context Length: %d", maxContextLen);

    bool ret = LlmManager::getInstance().initialize(nativeModelPath, maxNewTokens, maxContextLen, callback);
    // 释放 jstring 转换的内存
    env->ReleaseStringUTFChars(modelPath, nativeModelPath);

    return ret ? JNI_TRUE : JNI_FALSE;
}

extern "C" JNIEXPORT void  JNICALL
Java_com_taiic_deepseek_DsManager_run(
        JNIEnv* env,
        jobject /* this */,
        jstring input) {

    // 将 jstring 转换为 C++ 字符串
    const char* nativeInput = env->GetStringUTFChars(input, nullptr);
    if (nativeInput == nullptr) {
        // 内存分配失败
        return;
    }
    LlmManager::getInstance().run(nativeInput);
    // 释放 jstring 转换的内存
    env->ReleaseStringUTFChars(input, nativeInput);
}

extern "C" JNIEXPORT void  JNICALL
Java_com_taiic_deepseek_DsManager_destroy(
        JNIEnv* env,
        jobject /* this */) {

    g_javaVM->DetachCurrentThread(); // 释放线程与 JVM 的绑定

#if SIMPLE
    JavaSimpleManager::destoryJaveSimple(env);
#else
    JavaManager::destoryJavaMethod(env);
#endif

    LlmManager::getInstance().destroy();
}

extern "C" JNIEXPORT jboolean  JNICALL
Java_com_taiic_deepseek_DsManager_initialized(
        JNIEnv* env,
        jobject /* this */) {
    bool ret = LlmManager::getInstance().isInitialized();
    return ret ? JNI_TRUE : JNI_FALSE;
}